import numpy as np
import sensor_msgs.point_cloud2 as pc2


class LidarProcessor:
    def __init__(self, skip_points=2):
        """初始化激光雷达处理器"""
        self.skip_points = skip_points

    def process_pointcloud(self, cloud_msg):
        """处理点云数据"""
        points = []
        for p in pc2.read_points(
            cloud_msg, skip_points=self.skip_points, field_names=("x", "y", "z")
        ):
            points.append([p[0], p[1], p[2]])
        return np.array(points)

    def project_to_image(self, points_3d, camera_matrix, lidar_to_camera):
        """将点云投影到图像平面"""
        # 转换到相机坐标系
        points_camera = lidar_to_camera @ np.vstack(
            (points_3d.T, np.ones((1, points_3d.shape[0])))
        )

        # 投影到图像平面
        points_2d = camera_matrix @ points_camera[:3, :]
        points_2d = points_2d / points_2d[2, :]

        return points_2d[:2, :].T
